import onnxruntime
from transformers import BertTokenizer
import numpy as np
import torch
import torch.nn.functional as F


# 初始化ONNX Runtime会话
ort_session = onnxruntime.InferenceSession("bert_base_chinese.onnx")

# 加载分词器
tokenizer = BertTokenizer.from_pretrained(r"D:\pythonWork\python_demo\model\google-bert\bert-base-chinese\models--google-bert--bert-base-chinese\snapshots\c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f")

# 准备输入数据
text = "挺失望的,还不如买一本张爱玲文集呢,以<色戒>命名,可这篇文章仅仅10多页,且无头无尾的,完全比不上里面的任意一篇其它文章"
# text为输入的文本 return_tensors "np"(NumPy), "pt"(PyTorch), "tf"(TensorFlow) max_length 最大长度限制
inputs = tokenizer(text, padding="max_length", truncation=True, max_length=350, return_tensors="np")
# 运行推理
ort_inputs = {
    "input_ids": inputs["input_ids"].astype(np.int64),
    "attention_mask": inputs["attention_mask"].astype(np.int64),
    "token_type_ids": inputs["token_type_ids"].astype(np.int64)
}
ort_outputs = ort_session.run(None, ort_inputs)
#基本流程 获取输出->将输出转化为tensor->获取置信度概率->获取置信度概率最大的前 n 个结果->获取类别标签->处理输出
pred_logits = torch.tensor(ort_outputs[0])
# 对 logit 分数做 softmax 运算，得到置信度概率
pred_softmax = F.softmax(pred_logits, dim=1)
print(pred_softmax.shape)
# 获取置信度概率最大的前 n 个结果
n = 2
top_n = torch.topk(pred_softmax, n)
print(top_n)
pred_ids = top_n.indices.numpy()[0]
print(pred_ids)
confs = top_n.values.numpy()[0]
print(confs)
# 假设类别标签映射
id2label = {0: "负面评价", 1: "正面评价"}
print(id2label[pred_ids[0]])
# 处理输出
# logits = ort_outputs[0]
# predictions = np.argmax(logits, axis=-1)
# print(f"Predicted class: {predictions[0]}")